use buddy_system_allocator::LockedHeap;
use core::{ffi::c_void, mem::size_of};
use core::alloc::{
    GlobalAlloc,
    Layout
};

use crate::println;


#[global_allocator]
static HEAP_ALLOCATOR: LockedHeap<32> = LockedHeap::empty();

extern "C" {
    pub fn __bss_end();
}

const HEAP_END: usize = 0x60000000 + 8 * 1024 * 1024;

#[derive(Copy, Clone)]
struct AllocInfo {
    layout: Layout,
    ptr: *mut u8,
}

unsafe fn wrapped_alloc(layout: Layout) -> *mut u8 {
    let header_layout = Layout::new::<AllocInfo>();

    let (to_request, offset) = header_layout.extend(layout)
        .expect("real code should probably return null");

    let orig_ptr = HEAP_ALLOCATOR.alloc(to_request);
    if orig_ptr.is_null() {
        return orig_ptr;
    }

    let result_ptr = orig_ptr.add(offset);
    let info_ptr = result_ptr.sub(size_of::<AllocInfo>()) as *mut AllocInfo;
    info_ptr.write_unaligned(AllocInfo {
        layout: to_request,
        ptr: orig_ptr,
    });
    result_ptr
}

unsafe fn wrapped_dealloc(ptr: *mut u8) {
    assert!(!ptr.is_null());
    // Simply read the AllocInfo we wrote in `alloc`, and pass it into dealloc.
    let info_ptr = ptr.sub(size_of::<AllocInfo>()) as *const AllocInfo;
    let info = info_ptr.read_unaligned();
    HEAP_ALLOCATOR.dealloc(info.ptr, info.layout);
}

#[no_mangle]
pub fn init_heap() {
    extern "C" {
        pub fn _end();
    }
    println!("Init heap:0x{:x}~0x{:x}", __bss_end as usize, HEAP_END);
    unsafe {
        HEAP_ALLOCATOR
            .lock()
            .init(__bss_end as usize, HEAP_END - __bss_end as usize);
    }
}

#[no_mangle]
pub extern "C" fn rt_malloc(size: usize) -> *mut c_void {
    let addr;
    unsafe{
        addr = wrapped_alloc(Layout::from_size_align_unchecked(size, 4));
        // HEAP_ALLOCATOR.dealloc(ptr, layout)
    }
    addr as *mut c_void
}

#[no_mangle]
pub extern "C" fn rt_free(rmem: *mut c_void) {
    unsafe {
        wrapped_dealloc(rmem as *mut u8);
    }
}



#[alloc_error_handler]
pub fn handle_alloc_error(layout: core::alloc::Layout) -> ! {
    panic!("Heap allocation error, layout = {:?}", layout);
}